import tushare as ts
from typing import Dict, List, Optional
from datetime import datetime

class BasicPriceData:
    """基础价格数据类
    处理K线数据、分时数据、成交量数据和换手率等基础市场数据
    """
    
    def __init__(self, ts_token: str):
        """初始化基础价格数据类
        
        Args:
            ts_token: Tushare API token
        """
        self.ts_pro = ts.pro_api(ts_token)
        
    def get_daily_data(self, stock_code: str, start_date: str, end_date: str) -> Optional[Dict]:
        """获取日K线数据
        
        Args:
            stock_code: 股票代码
            start_date: 开始日期，格式YYYYMMDD
            end_date: 结束日期，格式YYYYMMDD
            
        Returns:
            包含开盘价、最高价、最低价、收盘价等数据的字典
        """
        try:
            df = self.ts_pro.daily(
                ts_code=stock_code,
                start_date=start_date,
                end_date=end_date
            )
            return df.to_dict('records') if not df.empty else None
        except Exception as e:
            print(f"获取日K线数据失败: {e}")
            return None
            
    def get_weekly_data(self, stock_code: str, start_date: str, end_date: str) -> Optional[Dict]:
        """获取周K线数据"""
        try:
            df = self.ts_pro.weekly(
                ts_code=stock_code,
                start_date=start_date,
                end_date=end_date
            )
            return df.to_dict('records') if not df.empty else None
        except Exception as e:
            print(f"获取周K线数据失败: {e}")
            return None
            
    def get_monthly_data(self, stock_code: str, start_date: str, end_date: str) -> Optional[Dict]:
        """获取月K线数据"""
        try:
            df = self.ts_pro.monthly(
                ts_code=stock_code,
                start_date=start_date,
                end_date=end_date
            )
            return df.to_dict('records') if not df.empty else None
        except Exception as e:
            print(f"获取月K线数据失败: {e}")
            return None
            
    def get_tick_data(self, stock_code: str, trade_date: str) -> Optional[Dict]:
        """获取分时数据"""
        try:
            df = ts.get_tick_data(stock_code, date=trade_date, src='tt')
            return df.to_dict('records') if not df.empty else None
        except Exception as e:
            print(f"获取分时数据失败: {e}")
            return None
            
    def get_turnover_rate(self, stock_code: str, trade_date: str) -> Optional[float]:
        """获取换手率数据"""
        try:
            df = self.ts_pro.daily_basic(
                ts_code=stock_code,
                trade_date=trade_date,
                fields='turnover_rate'
            )
            return float(df['turnover_rate'][0]) if not df.empty else None
        except Exception as e:
            print(f"获取换手率数据失败: {e}")
            return None